#include <define.h>

MODULE MOD_Block

!-------------------------------------------------------------------------------------
! !DESCRIPTION:
!
!    To deal with high-resolution data, the globe is divided into blocks.
!
!     (180W,90N)                           (180E,90N)
!        .-----------------------------------.
!        |         |         |        |      |
!        |         |         |        |      |
!        |         |         |        |      |
!        .-----------------------------------.
!        |         |         |        |      |
!        |         |         |        |      |
!        |         |         |        |      |
!        .-----------------------------------.
!        |         |         |        |      |
!        |         |         |        |      |
!        |         |         |        |      |
!        .-----------------------------------.
!     (180W,90S)                           (180E,90S)
!
!    1.
!    Boundaries for block (i,j) is saved in
!    "gblock%lat_s(j), gblock%lat_n(j), gblock%lon_w(i), gblock%lon_e(i)"
!    for south, north, west and east boundaries respectively.
!
!    2.
!    The (i,j) element of 2D array gblock%pio saves the global communication
!    number of process which is in charge of Input/Output of block (i,j).
!
!    3.
!    For Input/Output processes, "gblock%nblkme, gblock%xblkme(:), gblock%yblkme(:)"
!    SAVE the locations of blocks which are handled by themselves.
!
!    4.
!    Division of blocks can be generated by number of blocks globally (by set_by_size),
!    or set by predefined boundaries in files (by set_by_file).
!
!  Created by Shupeng Zhang, May 2023
!-------------------------------------------------------------------------------------

   USE MOD_Precision
   IMPLICIT NONE

   ! ---- data types ----
   type :: block_type

      ! Coordinates.
      integer :: nxblk, nyblk
      real(r8), allocatable :: lat_s (:)
      real(r8), allocatable :: lat_n (:)
      real(r8), allocatable :: lon_w (:)
      real(r8), allocatable :: lon_e (:)

      ! IO.
      integer, allocatable :: pio(:,:)

      integer :: nblkme
      integer, allocatable :: xblkme(:), yblkme(:)

   CONTAINS

      procedure, PUBLIC :: set => block_set

      procedure, PUBLIC :: save_to_file   => block_save_to_file
      procedure, PUBLIC :: load_from_file => block_load_from_file

      procedure, PRIVATE :: clip     => block_clip
      procedure, PRIVATE :: init_pio => block_init_pio
      procedure, PRIVATE :: read_pio => block_read_pio

      final :: block_free_mem

   END type block_type

   ! ---- Instance ----
   type (block_type) :: gblock


   ! ---- PUBLIC SUBROUTINE ----
   PUBLIC :: get_filename_block

CONTAINS

   ! --------------------------------
   SUBROUTINE block_set (this)

   USE MOD_Precision
   USE MOD_Namelist
   USE MOD_Utils
   USE MOD_SPMD_Task
   USE MOD_NetCDFSerial
   IMPLICIT NONE

   class (block_type) :: this

   ! Local Variables
   logical :: fexists
   integer :: iblk, jblk

      inquire(file=trim(DEF_BlockInfoFile), exist=fexists)

      IF (fexists) THEN

         CALL ncio_read_bcast_serial (DEF_BlockInfoFile, 'lat_s', this%lat_s)
         CALL ncio_read_bcast_serial (DEF_BlockInfoFile, 'lat_n', this%lat_n)
         CALL ncio_read_bcast_serial (DEF_BlockInfoFile, 'lon_w', this%lon_w)
         CALL ncio_read_bcast_serial (DEF_BlockInfoFile, 'lon_e', this%lon_e)

         this%nyblk = size(this%lat_s)
         this%nxblk = size(this%lon_w)

         ! blocks should be from south to north
         IF (this%lat_s(1) > this%lat_s(this%nyblk)) THEN
            this%lat_s = this%lat_s(this%nyblk:1:-1)
            this%lat_n = this%lat_n(this%nyblk:1:-1)
         ENDIF

      ELSE

         IF (DEF_AverageElementSize > 0) THEN

            this%nxblk = floor(360./(DEF_AverageElementSize/120.*50))
            this%nxblk = min(this%nxblk,360)
            DO WHILE ((this%nxblk < 360) .and. (mod(360,this%nxblk) /= 0))
               this%nxblk = this%nxblk + 1
            ENDDO

            this%nyblk = floor(180./(DEF_AverageElementSize/120.*50))
            this%nyblk = min(this%nxblk,180)
            DO WHILE ((this%nyblk < 180) .and. (mod(180,this%nyblk) /= 0))
               this%nyblk = this%nyblk + 1
            ENDDO

         ELSE

            this%nxblk = DEF_nx_blocks
            this%nyblk = DEf_ny_blocks

         ENDIF

         IF ((mod(360,DEF_nx_blocks) /= 0) .or. (mod(180,DEF_ny_blocks) /= 0)) THEN
            IF (p_is_master) THEN
               write(*,*) 'Number of blocks in longitude should be a factor of 360 '
               write(*,*) ' and Number of blocks in latitude should be a factor of 180.'
               CALL CoLM_stop ()
            ENDIF
         ENDIF

         allocate (this%lon_w (this%nxblk))
         allocate (this%lon_e (this%nxblk))

         DO iblk = 1, this%nxblk
            this%lon_w(iblk) = -180.0 + 360.0/this%nxblk * (iblk-1)
            this%lon_e(iblk) = -180.0 + 360.0/this%nxblk * iblk

            CALL normalize_longitude (this%lon_w(iblk))
            CALL normalize_longitude (this%lon_e(iblk))
         ENDDO

         allocate (this%lat_s (this%nyblk))
         allocate (this%lat_n (this%nyblk))

         DO jblk = 1, this%nyblk
            this%lat_s(jblk) = -90.0 + 180.0/this%nyblk * (jblk-1)
            this%lat_n(jblk) = -90.0 + 180.0/this%nyblk * jblk
         ENDDO

      ENDIF

#ifndef SinglePoint
      IF (p_is_master) THEN
         write (*,*)
         write (*,'(A)') '----- Block information -----'
         write (*,'(I4,A,I4,A)') this%nxblk, ' blocks in longitude,', &
            this%nyblk, ' blocks in latitude.'
         write (*,*)
      ENDIF
#else
      write(*,'(A)') 'Blocks : Set (360 longitude x 180 latitude) blocks for Single Point.'
#endif

#ifndef SinglePoint
      CALL this%init_pio ()
#endif

   END SUBROUTINE block_set

   ! --------------------------------
   SUBROUTINE block_save_to_file (this, dir_landdata)

   USE MOD_NetCDFSerial
   USE MOD_SPMD_Task
   IMPLICIT NONE

   class (block_type) :: this

   character(len=*), intent(in) :: dir_landdata

   ! Local variables
   character(len=256) :: filename

      IF (p_is_master) THEN

         filename = trim(dir_landdata) // '/block.nc'

         CALL ncio_create_file (filename)

         CALL ncio_define_dimension (filename, 'longitude', this%nxblk)
         CALL ncio_define_dimension (filename, 'latitude',  this%nyblk)

         CALL ncio_write_serial (filename, 'lat_s', this%lat_s, 'latitude' )
         CALL ncio_write_serial (filename, 'lat_n', this%lat_n, 'latitude' )
         CALL ncio_write_serial (filename, 'lon_w', this%lon_w, 'longitude')
         CALL ncio_write_serial (filename, 'lon_e', this%lon_e, 'longitude')

      ENDIF

   END SUBROUTINE block_save_to_file

   ! --------------------------------
   SUBROUTINE block_load_from_file (this, dir_landdata)

   USE MOD_NetCDFSerial
   USE MOD_SPMD_Task
   IMPLICIT NONE

   class (block_type) :: this
   character(len=*),  intent(in) :: dir_landdata

   ! Local variables
   character(len=256) :: filename

      filename = trim(dir_landdata) // '/block.nc'

      CALL ncio_read_bcast_serial (filename, 'lat_s', this%lat_s)
      CALL ncio_read_bcast_serial (filename, 'lat_n', this%lat_n)
      CALL ncio_read_bcast_serial (filename, 'lon_w', this%lon_w)
      CALL ncio_read_bcast_serial (filename, 'lon_e', this%lon_e)

      this%nyblk = size(this%lat_s)
      this%nxblk = size(this%lon_w)

      IF (p_is_master) THEN
         write (*,*) 'Block information:'
         write (*,'(I3,A,I3,A)') this%nxblk, ' blocks in longitude,', &
            this%nyblk, ' blocks in latitude.'
         write (*,*)
      ENDIF

      CALL this%read_pio (dir_landdata)

   END SUBROUTINE block_load_from_file

   ! --------------------------------
   SUBROUTINE block_clip (this, &
         iblk_south, iblk_north, iblk_west, iblk_east, numblocks)

   USE MOD_Namelist
   USE MOD_Utils
   IMPLICIT NONE

   class (block_type) :: this
   integer, intent(out) :: iblk_south, iblk_north, iblk_west, iblk_east
   integer, intent(out), optional :: numblocks

   ! Local Variables
   real(r8) :: edges, edgen, edgew, edgee
   integer  :: numblocks_x, numblocks_y

      edges = DEF_domain%edges
      edgen = DEF_domain%edgen
      edgew = DEF_domain%edgew
      edgee = DEF_domain%edgee

      iblk_south = find_nearest_south (edges, this%nyblk, this%lat_s)
      iblk_north = find_nearest_north (edgen, this%nyblk, this%lat_n)

      CALL normalize_longitude (edgew)
      CALL normalize_longitude (edgee)

      IF (edgew == edgee) THEN
         iblk_west = 1
         iblk_east = this%nxblk
      ELSE
         iblk_west = find_nearest_west (edgew, this%nxblk, this%lon_w)
         iblk_east = find_nearest_east (edgee, this%nxblk, this%lon_e)

         IF (iblk_west == iblk_east) THEN
            IF ((lon_between_floor(edgee,this%lon_w(iblk_west),edgew)) &
               .and. (this%lon_w(iblk_west) /= edgew)) THEN
               iblk_west = 1
               iblk_east = this%nxblk
            ENDIF
         ENDIF
      ENDIF

      IF (present(numblocks)) THEN

         numblocks_y = iblk_north - iblk_south + 1

         IF (iblk_east >= iblk_west) THEN
            numblocks_x = iblk_east - iblk_west + 1
         ELSE
            numblocks_x = this%nxblk - iblk_west + 1 + iblk_east
         ENDIF

         numblocks = numblocks_x * numblocks_y

      ENDIF

   END SUBROUTINE block_clip

   ! --------------------------------
   SUBROUTINE block_init_pio (this)

   USE MOD_Precision
   USE MOD_SPMD_Task
   USE MOD_Namelist
   USE MOD_Utils
   IMPLICIT NONE

   class (block_type) :: this

   integer :: iblk, jblk, iproc
   integer :: iblk_south, iblk_north, iblk_west, iblk_east
   integer :: numblocks, ngrp
   integer :: iblkme

      IF (p_is_master) THEN
         CALL this%clip (iblk_south, iblk_north, iblk_west, iblk_east, numblocks)
      ENDIF

#ifdef USEMPI
      CALL mpi_bcast (numblocks, 1, MPI_INTEGER, p_address_master, p_comm_glb, p_err)

      ngrp = max((p_np_glb-1) / DEF_PIO_groupsize, 1)
      ngrp = min(ngrp, numblocks)
      CALL divide_processes_into_groups (ngrp)
#endif

      allocate (this%pio (this%nxblk,this%nyblk))

      IF (p_is_master) THEN

         this%pio(:,:) = -1

         iproc = -1
         DO jblk = iblk_south, iblk_north

            iblk = iblk_west
            DO WHILE (.true.)
#ifdef USEMPI
               iproc = mod(iproc+1, p_np_io)
               this%pio(iblk,jblk) = p_address_io(iproc)
#else
               this%pio(iblk,jblk) = p_root
#endif

               IF (iblk /= iblk_east) THEN
                  iblk = mod(iblk,this%nxblk) + 1
               ELSE
                  EXIT
               ENDIF
            ENDDO
         ENDDO

      ENDIF

#ifdef USEMPI
      CALL mpi_bcast (this%pio, this%nxblk * this%nyblk, MPI_INTEGER, &
         p_address_master, p_comm_glb, p_err)
#endif

      this%nblkme = 0
      IF (p_is_io) THEN
         this%nblkme = count(this%pio == p_iam_glb)
         IF (this%nblkme > 0) THEN
            iblkme = 0
            allocate (this%xblkme(this%nblkme))
            allocate (this%yblkme(this%nblkme))
            DO iblk = 1, this%nxblk
               DO jblk = 1, this%nyblk
                  IF (p_iam_glb == this%pio(iblk,jblk)) THEN
                     iblkme = iblkme + 1
                     this%xblkme(iblkme) = iblk
                     this%yblkme(iblkme) = jblk
                  ENDIF
               ENDDO
            ENDDO
         ENDIF
      ENDIF

   END SUBROUTINE block_init_pio

   ! --------------------------------
   SUBROUTINE block_read_pio (this, dir_landdata)

   USE MOD_SPMD_Task
   USE MOD_NetCDFSerial
   USE MOD_Namelist
   IMPLICIT NONE

   class (block_type) :: this
   character(len=*),  intent(in) :: dir_landdata

   ! Local Variables
   character(len=256) :: filename, cyear
   integer, allocatable :: nelm_io(:), nelmblk(:,:)
   integer  :: iblk_south, iblk_north, iblk_west, iblk_east
   integer  :: numblocks, ngrp, iblk, jblk, iproc, jproc
   integer  :: iblkme

      IF (p_is_master) THEN
         ! Whether it varies by year???
         write(cyear,'(i4.4)') DEF_LC_YEAR
         filename = trim(dir_landdata) // '/mesh/' // trim(cyear) // '/mesh.nc'
         CALL ncio_read_serial (filename, 'nelm_blk', nelmblk)
         numblocks = count(nelmblk > 0)

         CALL this%clip (iblk_south, iblk_north, iblk_west, iblk_east)

      ENDIF

#ifdef USEMPI
      IF (p_is_master) THEN
         ngrp = max((p_np_glb-1) / DEF_PIO_groupsize, 1)
         ngrp = min(ngrp, numblocks)

         DO WHILE (.true.)

            allocate (nelm_io (ngrp))
            nelm_io(:) = 0

            DO jblk = iblk_south, iblk_north
               iblk = iblk_west
               DO WHILE (.true.)

                  IF (nelmblk(iblk,jblk) > 0) THEN
                     iproc = minloc(nelm_io, dim=1)
                     nelm_io(iproc) = nelm_io(iproc) + nelmblk(iblk,jblk)
                  ENDIF

                  IF (iblk /= iblk_east) THEN
                     iblk = mod(iblk,this%nxblk) + 1
                  ELSE
                     EXIT
                  ENDIF
               ENDDO
            ENDDO

            IF (maxval(nelm_io) < 2 * minval(nelm_io)) THEN
               deallocate (nelm_io)
               EXIT
            ELSE
               ngrp = ngrp - 1
               deallocate (nelm_io)
            ENDIF
         ENDDO
      ENDIF

      CALL mpi_bcast (numblocks, 1, MPI_INTEGER, p_address_master, p_comm_glb, p_err)
      CALL mpi_bcast (ngrp,      1, MPI_INTEGER, p_address_master, p_comm_glb, p_err)
      CALL divide_processes_into_groups (ngrp)
#endif

      allocate (this%pio (this%nxblk,this%nyblk))

      IF (p_is_master) THEN

         this%pio(:,:) = -1

#ifdef USEMPI
         allocate (nelm_io (0:p_np_io-1))
         nelm_io(:) = 0
         jproc = -1
#endif

         DO jblk = iblk_south, iblk_north
            iblk = iblk_west
            DO WHILE (.true.)
#ifdef USEMPI
               IF (nelmblk(iblk,jblk) > 0) THEN
                  iproc = minloc(nelm_io, dim=1) - 1
                  this%pio(iblk,jblk) = p_address_io(iproc)
                  nelm_io(iproc) = nelm_io(iproc) + nelmblk(iblk,jblk)
               ELSEIF (nelmblk(iblk,jblk) == 0) THEN
                  jproc = mod(jproc+1, p_np_io)
                  this%pio(iblk,jblk) = p_address_io(jproc)
               ENDIF
#else
               this%pio(iblk,jblk) = p_root
#endif

               IF (iblk /= iblk_east) THEN
                  iblk = mod(iblk,this%nxblk) + 1
               ELSE
                  EXIT
               ENDIF
            ENDDO
         ENDDO

#ifdef USEMPI
         deallocate (nelm_io)
#endif
      ENDIF

#ifdef USEMPI
      CALL mpi_bcast (this%pio, this%nxblk * this%nyblk, MPI_INTEGER, &
         p_address_master, p_comm_glb, p_err)
#endif

      this%nblkme = 0
      IF (p_is_io) THEN
         this%nblkme = count(this%pio == p_iam_glb)
         IF (this%nblkme > 0) THEN
            iblkme = 0
            allocate (this%xblkme(this%nblkme))
            allocate (this%yblkme(this%nblkme))
            DO iblk = 1, this%nxblk
               DO jblk = 1, this%nyblk
                  IF (p_iam_glb == this%pio(iblk,jblk)) THEN
                     iblkme = iblkme + 1
                     this%xblkme(iblkme) = iblk
                     this%yblkme(iblkme) = jblk
                  ENDIF
               ENDDO
            ENDDO
         ENDIF
      ENDIF

      IF (allocated(nelmblk)) deallocate (nelmblk)

   END SUBROUTINE block_read_pio

   ! --------------------------------
   SUBROUTINE block_free_mem (this)

   IMPLICIT NONE
   type (block_type) :: this

      IF (allocated (this%lat_s))  deallocate (this%lat_s)
      IF (allocated (this%lat_n))  deallocate (this%lat_n)
      IF (allocated (this%lon_w))  deallocate (this%lon_w)
      IF (allocated (this%lon_e))  deallocate (this%lon_e)

      IF (allocated (this%pio)  )  deallocate (this%pio  )

      IF (allocated (this%xblkme)) deallocate (this%xblkme)
      IF (allocated (this%yblkme)) deallocate (this%yblkme)

   END SUBROUTINE block_free_mem

   ! -----
   SUBROUTINE get_blockname (iblk, jblk, blockname)

   IMPLICIT NONE

   integer, intent(in) :: iblk, jblk

   character(len=*), intent(out) :: blockname

   ! Local variables
   character(len=4) :: cx
   character(len=3) :: cy
   integer :: i

      IF (gblock%lat_s(jblk) < 0) THEN
         write (cy, '(A1,I2.2)') 's', - floor(gblock%lat_s(jblk))
      ELSE
         write (cy, '(A1,I2.2)') 'n',   floor(gblock%lat_s(jblk))
      ENDIF

      IF (gblock%lon_w(iblk) < 0) THEN
         write (cx, '(A1,I3.3)') 'w', - floor(gblock%lon_w(iblk))
      ELSE
         write (cx, '(A1,I3.3)') 'e',   floor(gblock%lon_w(iblk))
      ENDIF

      blockname = trim(cx) // '_' // trim(cy)

   END SUBROUTINE get_blockname

   ! --------------------------------
   SUBROUTINE get_filename_block (filename, iblk, jblk, fileblock)

   IMPLICIT NONE

   character(len=*), intent(in) :: filename
   integer, intent(in) :: iblk, jblk

   character(len=*), intent(out) :: fileblock

   ! Local variables
   character(len=8) :: blockname
   integer :: i

      CALL get_blockname (iblk, jblk, blockname)

      i = len_trim (filename)
      DO WHILE (i > 0)
         IF (filename(i:i) == '.') EXIT
         i = i - 1
      ENDDO

      IF (i > 0) THEN
         fileblock = filename(1:i-1) // '_' // blockname // '.nc'
      ELSE
         fileblock = filename // '_' // blockname // '.nc'
      ENDIF

   END SUBROUTINE get_filename_block

END MODULE MOD_Block
